#!/bin/bash -eux
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
################################################################################

# In order to enable PySecSan for a given module, set the environment
# variable ENABLE_PYSECSAN="YES"

fuzzer_path=$1
shift 1

fuzzer_basename=$(basename -s .py $fuzzer_path)
fuzzer_package=${fuzzer_basename}.pkg

PYFUZZ_WORKPATH=$SRC/pyfuzzworkdir/
FUZZ_WORKPATH=$PYFUZZ_WORKPATH/$fuzzer_basename

if [[ $SANITIZER = *introspector* ]]; then
    # Extract the source package the fuzzer targets. This must happen before
    # we enter the virtual environment in the following lines because we need
    # to use the same python environment that installed the fuzzer dependencies.
    python3 /fuzz-introspector/frontends/python/prepare_fuzz_imports.py $fuzzer_path isossfuzz

    # We must ensure python3.9, this is because we use certain
    # AST logic from there.
    # The below should probably be refined
    apt-get install -y python3.9
    apt-get update
    apt-get install -y python3-pip
    python3.9 -m pip install virtualenv
    python3.9 -m virtualenv .venv
    . .venv/bin/activate
    pip3 install pyyaml
    export PYTHONPATH="/fuzz-introspector/frontends/python/PyCG"

    ARGS="--fuzzer $fuzzer_path"
    if [ -n "${PYFUZZPACKAGE-}" ]; then
      ARGS="$ARGS --package=${PYFUZZPACKAGE}"
    fi
    python /fuzz-introspector/frontends/python/main.py $ARGS
    ls -la ./
    exit 0
fi

# In coverage mode prepend coverage logic to the fuzzer source
if [[ $SANITIZER = *coverage* ]]; then
  cat <<EOF > coverage_wrapper.py
###### Coverage stub
import atexit
import coverage
cov = coverage.coverage(data_file='.coverage', cover_pylib=True)
cov.start()
# Register an exist handler that will print coverage
def exit_handler():
    cov.stop()
    cov.save()
atexit.register(exit_handler)
####### End of coverage stub
EOF

  # Prepend stub and create tmp file
  cat coverage_wrapper.py $fuzzer_path > tmp_fuzzer_coverage.py

  # Overwrite existing fuzzer with new fuzzer that has stub
  mv tmp_fuzzer_coverage.py $fuzzer_path
fi

# If PYSECSAN is enabled, ensure that we can build with it.
if [[ ${ENABLE_PYSECSAN:-"0"} != "0" ]];
then
  # Make sure pysecsan is installed
  if [[ ! -d "/pysecsan" ]];
  then
    pushd /usr/local/lib/sanitizers/pysecsan
    python3 -m pip install .
    popd
  fi

  cat <<EOF > pysecsan_wrapper.py
import pysecsan; pysecsan.add_hooks();
EOF

  # Prepend stub and create tmp file
  cat pysecsan_wrapper.py $fuzzer_path > tmp_fuzzer_pysecsan.py

  # Overwrite existing fuzzer with new fuzzer that has stub
  mv tmp_fuzzer_pysecsan.py $fuzzer_path
fi

rm -rf $PYFUZZ_WORKPATH
mkdir $PYFUZZ_WORKPATH $FUZZ_WORKPATH

pyinstaller --distpath $OUT --workpath=$FUZZ_WORKPATH --onefile --name $fuzzer_package "$@" $fuzzer_path

# Disable executable bit from package as OSS-Fuzz uses executable bits to
# identify fuzz targets. We re-enable the executable bit in wrapper script
# below.
chmod -x $OUT/$fuzzer_package

# In coverage mode save source files of dependencies in pyinstalled binary
if [[ $SANITIZER = *coverage* ]]; then
  rm -rf /medio/
  python3 /usr/local/bin/python_coverage_helper.py $FUZZ_WORKPATH "/medio"
  zip -r $fuzzer_package.deps.zip /medio
  mv $fuzzer_package.deps.zip $OUT/
fi

# Create execution wrapper.
echo "#!/bin/sh
# LLVMFuzzerTestOneInput for fuzzer detection.
this_dir=\$(dirname \"\$0\")
chmod +x \$this_dir/$fuzzer_package
LD_PRELOAD=\$this_dir/sanitizer_with_fuzzer.so \
ASAN_OPTIONS=\$ASAN_OPTIONS:symbolize=1:external_symbolizer_path=\$this_dir/llvm-symbolizer:detect_leaks=0 \
\$this_dir/$fuzzer_package \$@" > $OUT/$fuzzer_basename
chmod +x $OUT/$fuzzer_basename
